#!/usr/bin/env python2.7

# bender.py

import os, sys
#import re
#from datetime import datetime, timedelta
#from dateutil.parser import parse
import numpy as np
import matplotlib
matplotlib.use('Qt4Agg')
matplotlib.rcParams['backend.qt4']='PySide'
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
from matplotlib.figure import Figure

#from mpl_toolkits.basemap import Basemap
#from matplotlib import mlab
#from netCDF4 import Dataset
from PySide.QtCore import *
from PySide.QtGui import *
from ui_lidarimage import *
from metlib.kits import Null, isseq
import metlib.color
from metlib.lidar import *
from metlib.plot.artist_util import *

cmapd = dict()

cmapnames = []
for m in matplotlib.cm.datad:
    if not m.endswith("_r"):
        cmapd[m] = matplotlib.cm.__dict__[m]
        cmapd[m+'_r'] = matplotlib.cm.__dict__[m+'_r']
        cmapnames.append(m)

for m in metlib.color.cm.__all__:
    if m.startswith('cm_') and not m.endswith("_r"):
        cmapd[m] = metlib.color.cm.__dict__[m]
        cmapd[m+'_r'] = metlib.color.cm.__dict__[m+'_r']
        cmapnames.append(m)

cmapnames.sort()

class LidarImageBase(QWidget, Ui_LidarImageWidget):
    def __init__(self, parent=None):
        super(LidarImageBase, self).__init__(parent)
        self.setupUi(self)
        self.fig = Figure(figsize=(12,5.5), dpi=96)
        self.ax = self.fig.add_axes((0.06,0.08,0.82,0.80))
        self.ax.autoscale(False)
        #self.cax = self.fig.add_axes((0.91, 0.12, 0.03, 0.82))
        self.canvas = FigureCanvas(self.fig)
        self.canvas.setFocusPolicy( Qt.ClickFocus )
        self.imageLayout.addWidget(self.canvas)
        self.mpl_toolbar = NavigationToolbar(self.canvas, self)
        self.toolbarLayout.addWidget(self.mpl_toolbar)
        self.cmapBox.addItems(cmapnames)
        self.cmapBox.setCurrentIndex(self.cmapBox.findText('cm_night'))
        self.cmapBox.currentIndexChanged.connect(self.update_cmap)
        self.cmap_r_checkBox.stateChanged.connect(self.update_cmap)
        self.drawButton.clicked.connect(self.update_cmap)
    
        self.extra_data = []
        self.extra_artist_ids = ArtistIDManager()
        self.index2id = []

    def update_fig(self, *args):
        self.canvas.draw()

    def update_cmap(self, *args):
        pass

    def load_data(self, *args):
        self.clear_extra_data()

    def clear_extra_data(self, i=None):
        if i is None:
            del self.extra_data[0:len(self.extra_data)]
            del self.index2id[0:len(self.index2id)]
            self.extra_artist_ids.clear()
        else:
            del self.extra_data[i]
            del self.extra_artist_ids[self.index2id[i]]
            del self.index2id[i]
        self.canvas.draw()

    def clear_extra_data_by_id(self, IDs=[]):
        if not isseq(IDs):
            IDs = [IDs]
        for ID in IDs:
            try:
                i = self.index2id.index(ID)
                self.clear_extra_data(i)
            except ValueError as e:
                pass

    def clear_extra_last(self):
        if self.extra_data:
            self.extra_data.pop()
            del self.extra_artist_ids[self.index2id[-1]]
            del self.index2id[-1]
            self.canvas.draw()

    def load_extra_data(self, plottype, dts, heights, *args, **kwargs):
        self.extra_data.append([plottype, dts, heights, args, kwargs])

        if 'mark' in kwargs:
            del kwargs['mark']

        if 'ID' in kwargs:
            newID = kwargs['ID']
            del kwargs['ID']
        else:
            newID = self.extra_artist_ids.get_new_ID()
        if plottype == 'plot':
            artist = self.ax.plot(dts, heights, *args, **kwargs)[0]
        elif plottype == 'scatter':
            artist = self.ax.scatter(dts, heights, *args, **kwargs)
        elif plottype == 'text':
            artist = self.ax.text(dts, heights, *args, **kwargs)
        self.extra_artist_ids[newID] = artist
        self.index2id.append(newID)
#        try:
#            self.ax.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc='lower left', ncol=4)
#        except Exception as e:
#            sys.stderr.write(e)
#            sys.stderr.write('\n')
        self.canvas.draw()

    @property
    def cmap(self):
        return cmapd[self.cmapname]

    @property
    def cmapname(self):
        appendix = '_r' if self.cmap_r_checkBox.isChecked() else ''
        return self.cmapBox.currentText()+appendix

    @cmapname.setter
    def cmapname(self, value):
        if value.endswith('_r'):
            value = value.rstrip('_r')
            reverse = Qt.Checked
        else:
            reverse = Qt.Unchecked
        self.cmapBox.setCurrentIndex(self.cmapBox.findText(value))
        self.cmap_r_checkBox.setCheckState(reverse)

    @property
    def vmin(self):
        return self.vminSpinBox.value()

    @vmin.setter
    def vmin(self, value):
        self.vminSpinBox.setValue(value)
    
    @property
    def vmax(self):
        return self.vmaxSpinBox.value()

    @vmax.setter
    def vmax(self, value):
        self.vmaxSpinBox.setValue(value)

    @property
    def vlock(self):
        return self.vlock_checkBox.isChecked()

    @vlock.setter
    def vlock(self, value):
        state = Qt.Checked if value else Qt.Unchecked
        self.vlock_checkBox.setCheckState(state)

class LidarImage(LidarImageBase):
    def __init__(self, parent=None):
        super(LidarImage, self).__init__(parent)
        self.cax = self.fig.add_axes((0.91, 0.08, 0.03, 0.80))
        self.image = Null

    def load_data(self, dts, heights, data):
        self.clear_extra_data()
        self.dts = dts
        self.heights = heights
        self.data = data
    
    def update_fig(self, *args):
        self.image.remove()
        self.cax.clear()
        self.image = plot_lidar_on_ax(self.dts, self.heights, self.data,
                self.ax, cax=self.cax,
                vmin=self.vmin, vmax=self.vmax,
                cmap=self.cmap,
                plotter='imshow',
                zorder=0
                )
        self.canvas.draw()
   
    def update_cmap(self):
        self.image.set_cmap(self.cmap)
        self.image.set_clim(self.vmin, self.vmax)
        self.canvas.draw()

class LidarSketchImage(LidarImage):
    lineDrawn = Signal(object)

    def __init__(self, parent=None):
        super(LidarSketchImage, self).__init__(parent)
        
        self.colord = {1:'k', 2:'b', 3:'c', 4:'m', 5:'r', 6:'Orange', 7:'y', 8:'g'}
        self.ci = 1
        self.color = self.colord[self.ci]
        self.nowID = 1
        self.xs = []
        self.ys = []
        self.tmp_points = []
        self.pressed = False

        self.canvas.mpl_connect('button_press_event', self.on_press)
        self.canvas.mpl_connect('motion_notify_event', self.on_motion)
        self.canvas.mpl_connect('button_release_event', self.on_release)

    def load_data(self, dts, heights, data):
        self.reset_ID()
        super(LidarSketchImage, self).load_data(dts, heights, data)

    def set_color(self, ci):
        self.ci = ci
        self.color = self.colord[ci]

    def reset_ID(self):
        self.nowID = 1

    def on_press(self, event):
        if self.mpl_toolbar.mode != '':
            return
        if event.inaxes != self.ax: 
            return
        if event.button != 1:
            return
        for tmp_point in self.tmp_points:
            tmp_point.remove()
        self.xs = [event.xdata]
        self.ys = [event.ydata]
        self.tmp_points = [self.ax.scatter(event.xdata, event.ydata, c=self.color, s=20, alpha=0.5, linewidths=0.0)]
        self.pressed = True

        self.canvas.draw()

    def on_motion(self, event):
        if self.mpl_toolbar.mode != '':
            return
        if event.inaxes != self.ax: 
            return
        if event.button != 1:
            return
        if self.pressed == False:
            return
        self.xs.append(event.xdata)
        self.ys.append(event.ydata)
        self.tmp_points.append(self.ax.scatter(event.xdata, event.ydata, c=self.color, s=20, alpha=0.5, linewidths=0.0))
        self.canvas.draw()

    def on_release(self, event):
        if self.mpl_toolbar.mode != '':
            return
        if event.button != 1:
            return
        self.pressed = False
        for tmp_point in self.tmp_points:
            tmp_point.remove()
        self.tmp_points = []
        if len(self.xs) == 0:
            return
        lineID = 'L%d:%d' % (self.ci, self.nowID)
        self.draw_sketchline(lineID, self.xs, self.ys)
        self.xs = []
        self.ys = []

    def draw_sketchline(self, lineID, xs, ys):
        tks = lineID.lstrip('L').split(':')
        cid = int(tks[0])
        iid = int(tks[1])
        color = self.colord[cid]
        self.load_extra_data('plot', xs, ys, color=color, alpha=0.7, linewidth=2.5, ID=lineID, mark='sketchline')
        self.load_extra_data('text', xs[0], ys[0], lineID, fontsize='x-small', fontweight='semibold', color='k', alpha=1.0, ha='right', va='bottom', ID=lineID)
        self.lineDrawn.emit((lineID, xs, ys))
        self.canvas.draw()
        if self.nowID <= iid:
            self.nowID = iid+1
        

class LidarPeakImage(LidarImageBase):
    def __init__(self, parent=None, ref_image=None):
        super(LidarPeakImage, self).__init__(parent)
        self.cmapBox.setEnabled(False)
        self.cmap_r_checkBox.setEnabled(False)
        self.peak_scatter = Null
        self.ref_image = ref_image
        self.canvas.mpl_connect('pick_event', self.on_pick)
        self.sel_ij = -1, -1
        self.sel_line = None
        self.sel_line_i = None

    def load_data(self, dts, heights, ss, cs):
        self.dts = dts
        self.heights = heights
        self.ss = ss
        self.cs = cs
        self.clear_extra_data()

    def update_fig(self, *args, **kwargs):
        self.peak_scatter.remove()
        self.peak_scatter = self.ax.scatter(self.dts, self.heights, self.ss, self.cs, linewidth=0, picker=5, zorder=0)
        self.set_xylim()
        self.canvas.draw()

    def set_xylim(self):
        if self.ref_image is not None:
            self.ax.set_xlim(self.ref_image.ax.get_xlim())
            self.ax.set_ylim(self.ref_image.ax.get_ylim())

    def on_pick(self, event):
        sc_ind = event.ind[len(event.ind)/2]
        if event.mouseevent.button == 1:
            pointtype = 'BumpPoint'
            if event.artist in self.extra_artist_ids:
                return
            else:
                self.sel_line = None
                self.sel_line_i = None
                self.sel_t = self.dts[sc_ind]
                self.sel_h = self.heights[sc_ind]
        elif event.mouseevent.button == 3:
            pointtype = 'TrimPoint'
            if not event.artist in self.extra_artist_ids:
                return
            else:
                try:
                    self.sel_line = event.artist
                    self.sel_line_i = sc_ind
                    line_id = self.extra_artist_ids[self.sel_line]
                    extra_i = self.index2id.index(line_id)
                    self.sel_t = self.extra_data[extra_i][1][sc_ind]
                    self.sel_h = self.extra_data[extra_i][2][sc_ind]
                except Exception as e:
                    print self.sel_line
                    print self.sel_line_i
                    print self.extra_artist_ids.ID2artist
                    print line_id
                    print self.index2id
                    print extra_i
                    print len(self.extra_data)

        self.info_Label.setText('%s %s %.1f' % (pointtype, self.sel_t, self.sel_h))


if __name__ == '__main__':
    pass
